"""
Emittance calculation utilities for beam analysis
Calculates RMS, geometric, and normalized parameters from particle data
"""

import numpy as np
import pandas as pd
from scipy.spatial import ConvexHull
from .convex_hanlder import calculate_convex_hull_area


def calculate_rms_emittance(x, x_prime):
    """
    Calculate RMS emittance using the correct formula from SIMION documentation
    ε_rms = √[⟨x²⟩⟨x'²⟩ - ⟨xx'⟩²]
    
    Args:
        x (array): Position coordinates
        x_prime (array): Slope coordinates (x' = vx/vz)
        
    Returns:
        float: RMS emittance
    """
    # Remove means to center the phase space ellipse
    x_centered = x - np.mean(x)
    x_prime_centered = x_prime - np.mean(x_prime)
    
    # Calculate second order moments
    # ⟨x²⟩ = variance of x
    sigma_11 = np.mean(x_centered**2)
    # ⟨xx'⟩ = covariance of x and x'
    sigma_12 = np.mean(x_centered * x_prime_centered)
    # ⟨x'²⟩ = variance of x'
    sigma_22 = np.mean(x_prime_centered**2)
    
    # Calculate RMS emittance: ε_rms = √[⟨x²⟩⟨x'²⟩ - ⟨xx'⟩²]
    emittance = np.sqrt(sigma_11 * sigma_22 - sigma_12**2)
    
    return emittance


def calculate_rms_parameters(df):
    """
    Calculate RMS parameters for beam analysis using correct emittance formula
    
    Args:
        df (DataFrame): DataFrame with columns ['x', 'y', 'vx', 'vy', 'vz']
        
    Returns:
        dict: Dictionary containing RMS parameters
    """
    x = df['x'].values
    y = df['y'].values
    vx = df['vx'].values
    vy = df['vy'].values
    vz = df['vz'].values
    
    # Calculate slopes in trace space (x' = vx/vz, y' = vy/vz)
    vz_nonzero = np.where(vz != 0, vz, 1e-10)
    x_prime = vx / vz_nonzero
    y_prime = vy / vz_nonzero
    
    # Calculate RMS emittance for x and y directions
    rms_emittance_x = calculate_rms_emittance(x, x_prime)
    rms_emittance_y = calculate_rms_emittance(y, y_prime)
    
    # Calculate geometric mean of RMS emittances
    rms_gm_xy = np.sqrt(rms_emittance_x * rms_emittance_y)
    
    # Also calculate traditional RMS values for reference
    rms_x = np.sqrt(np.mean(x**2))
    rms_y = np.sqrt(np.mean(y**2))
    rms_vx = np.sqrt(np.mean(vx**2))
    rms_vy = np.sqrt(np.mean(vy**2))
    rms_vz = np.sqrt(np.mean(vz**2))
    
    return {
        'rms_x': rms_x,
        'rms_y': rms_y,
        'rms_vx': rms_vx,
        'rms_vy': rms_vy,
        'rms_vz': rms_vz,
        'rms_emittance_x': rms_emittance_x,
        'rms_emittance_y': rms_emittance_y,
        'rms_gm_xy': rms_gm_xy
    }


def calculate_geometric_emittance(rms_emittance):
    """
    Calculate geometric emittance from RMS emittance
    Formula: ε_x = 4 ε_rms,x
    
    Args:
        rms_emittance (float): RMS emittance value
        
    Returns:
        float: Geometric emittance
    """
    return 4.0 * rms_emittance


def calculate_normalized_emittance(geometric_emittance, velocity):
    """
    Calculate normalized emittance
    Formula: ε_n = βγ ε
    where β = v/c and γ = 1/√(1 - β²)
    
    Args:
        geometric_emittance (float): Geometric emittance
        velocity (float): Average velocity of particles (m/s)
        
    Returns:
        float: Normalized emittance
    """
    c = 299792458.0  # Speed of light in m/s
    beta = velocity / c
    gamma = 1.0 / np.sqrt(1.0 - beta**2)
    
    return beta * gamma * geometric_emittance


def calculate_geometric_parameters(df):
    """
    Calculate geometric parameters using RMS emittance
    Formula: ε_x = 4 ε_rms,x
    
    Args:
        df (DataFrame): DataFrame with columns ['x', 'y', 'vx', 'vy', 'vz']
        
    Returns:
        dict: Dictionary containing geometric parameters
    """
    # First calculate RMS parameters to get RMS emittances
    rms_params = calculate_rms_parameters(df)
    
    # Calculate geometric emittances using the formula ε_x = 4 ε_rms,x
    geometric_emittance_x = calculate_geometric_emittance(rms_params['rms_emittance_x'])
    geometric_emittance_y = calculate_geometric_emittance(rms_params['rms_emittance_y'])
    
    # Calculate geometric mean
    geometric_gm_xy = np.sqrt(geometric_emittance_x * geometric_emittance_y)
    
    return {
        'geometric_x': geometric_emittance_x,
        'geometric_y': geometric_emittance_y,
        'geometric_gm_xy': geometric_gm_xy
    }


def calculate_normalized_parameters(df):
    """
    Calculate normalized parameters using the correct formula
    Formula: ε_n = βγ ε
    where β = v/c and γ = 1/√(1 - β²)
    
    Args:
        df (DataFrame): DataFrame with columns ['x', 'y', 'vx', 'vy', 'vz']
        
    Returns:
        dict: Dictionary containing normalized parameters
    """
    # First calculate RMS and geometric parameters
    rms_params = calculate_rms_parameters(df)
    geometric_params = calculate_geometric_parameters(df)
    
    # Calculate average velocity magnitude
    vx = df['vx'].values
    vy = df['vy'].values
    vz = df['vz'].values
    velocity_magnitude = np.sqrt(vx**2 + vy**2 + vz**2)
    average_velocity = np.mean(velocity_magnitude)
    
    # Calculate normalized emittances using the formula ε_n = βγ ε
    normalized_x = calculate_normalized_emittance(geometric_params['geometric_x'], average_velocity)
    normalized_y = calculate_normalized_emittance(geometric_params['geometric_y'], average_velocity)
    
    # Calculate geometric mean
    normalized_gm_xy = np.sqrt(normalized_x * normalized_y)
    
    return {
        'normalized_x': normalized_x,
        'normalized_y': normalized_y,
        'normalized_gm_xy': normalized_gm_xy,
        'average_velocity': average_velocity
    }


def calculate_all_parameters(df):
    """
    Calculate all beam parameters (RMS, geometric, normalized)
    
    Args:
        df (DataFrame): DataFrame with columns ['x', 'y', 'vx', 'vy', 'vz']
        
    Returns:
        dict: Dictionary containing all parameters
    """
    rms_params = calculate_rms_parameters(df)
    geometric_params = calculate_geometric_parameters(df)
    normalized_params = calculate_normalized_parameters(df)
    
    # Combine all parameters
    all_params = {}
    all_params.update(rms_params)
    all_params.update(geometric_params)
    all_params.update(normalized_params)
    
    return all_params


def process_files_for_emittance_analysis(dfs, filenames=None):
    """
    Process multiple files to extract emittance parameters vs frequency
    
    Args:
        dfs (list): List of DataFrames
        filenames (list): List of filenames (optional)
        
    Returns:
        dict: Dictionary containing frequency and all parameters
    """
    frequencies = []
    rms_emittance_x_list = []
    rms_emittance_y_list = []
    rms_gm_xy_list = []
    geometric_x_list = []
    geometric_y_list = []
    geometric_gm_xy_list = []
    normalized_x_list = []
    normalized_y_list = []
    normalized_gm_xy_list = []
    area_list = []

    for i, df in enumerate(dfs):
        try:
            # Extract frequency directly from DataFrame 'f' column
            freq = df['f'].iloc[0]
            
            # Calculate all parameters
            params = calculate_all_parameters(df)
            
            # Calculate convex hull area
            area = calculate_convex_hull_area(df['x'], df['y'])
            
            # Store results
            frequencies.append(freq)
            rms_emittance_x_list.append(params['rms_emittance_x'])
            rms_emittance_y_list.append(params['rms_emittance_y'])
            rms_gm_xy_list.append(params['rms_gm_xy'])
            geometric_x_list.append(params['geometric_x'])
            geometric_y_list.append(params['geometric_y'])
            geometric_gm_xy_list.append(params['geometric_gm_xy'])
            normalized_x_list.append(params['normalized_x'])
            normalized_y_list.append(params['normalized_y'])
            normalized_gm_xy_list.append(params['normalized_gm_xy'])
            area_list.append(area)
            
        except Exception as e:
            print(f"Error processing file {i+1}: {str(e)}")
            # Add zeros for failed files
            frequencies.append(i)
            rms_emittance_x_list.append(0)
            rms_emittance_y_list.append(0)
            rms_gm_xy_list.append(0)
            geometric_x_list.append(0)
            geometric_y_list.append(0)
            geometric_gm_xy_list.append(0)
            normalized_x_list.append(0)
            normalized_y_list.append(0)
            normalized_gm_xy_list.append(0)
            area_list.append(0)
    
    return {
        'frequencies': np.array(frequencies),
        'rms_emittance_x': np.array(rms_emittance_x_list),
        'rms_emittance_y': np.array(rms_emittance_y_list),
        'rms_gm_xy': np.array(rms_gm_xy_list),
        'geometric_x': np.array(geometric_x_list),
        'geometric_y': np.array(geometric_y_list),
        'geometric_gm_xy': np.array(geometric_gm_xy_list),
        'normalized_x': np.array(normalized_x_list),
        'normalized_y': np.array(normalized_y_list),
        'normalized_gm_xy': np.array(normalized_gm_xy_list),
        'area': np.array(area_list)
    }
